import torch
from models.net import LeNet
from torchvision.transforms import Compose, ToTensor, Grayscale, Resize
from PIL import Image
import os

# 定义图像预处理
data_transform = Compose([
    Grayscale(),  # 转为灰度图
    Resize((28, 28)),  # 调整为28x28像素
    ToTensor()  # 转为Tensor格式
])

# 如果有NVIDIA显卡，转到GPU训练，否则用CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 模型实例化，将模型转到device
model = LeNet().to(device)

# 加载训练好的模型
model.load_state_dict(torch.load("save_model/best_model.pth"))

# 结果类型
classes = [
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
]

# 定义大文件夹路径
root_folder = "Tmage"  # 替换为大文件夹路径

# 切换模型为评估模式
model.eval()

# 初始化计数器
total_images = 0
correct_predictions = 0

# 遍历每个子文件夹
for folder_name in os.listdir(root_folder):
    folder_path = os.path.join(root_folder, folder_name)
    if not os.path.isdir(folder_path) or not folder_name.isdigit():  # 确保是数字命名的文件夹
        continue

    true_label = folder_name  # 文件夹名即为正确标签

    # 遍历该文件夹中的所有图片
    for filename in os.listdir(folder_path):
        image_path = os.path.join(folder_path, filename)
        if not filename.lower().endswith(('.png', '.jpg')):  # 过滤非图片文件
            continue

        # 加载并预处理图片
        image = Image.open(image_path)
        image = data_transform(image).unsqueeze(0).to(device)

        # 预测结果
        with torch.no_grad():
            pred = model(image)
            predicted_label = classes[torch.argmax(pred[0])]

        print(f'File:{folder_name}: {filename}, Predicted: "{predicted_label}"')

        # 统计总图片数和正确预测数
        total_images += 1
        if predicted_label == true_label:
            correct_predictions += 1

# 计算正确率
accuracy = (correct_predictions / total_images) * 100 if total_images > 0 else 0
print(f"Total Images: {total_images}, Correct Predictions: {correct_predictions}, Accuracy: {accuracy:.2f}%")